#include <assert.h>
#include <stdio.h>

#include "../src/mstep.c"

static void
fill (double*           x,
      double          (*func) (void),
      unsigned int      n)
{
  for (unsigned int i = 0; i < n; ++i)
    {
      x[i] = func ();
    }
}

static void
fill_r (Rating*         x,
        unsigned int    n_workers,
        unsigned int    n_labels,
        unsigned int    n_images,
        unsigned int    n)
{
  for (unsigned int i = 0; i < n; ++i)
    {
      x[i].worker = (lrand48 () >> 16) % n_workers;
      x[i].image = (lrand48 () >> 16) % n_images;
      x[i].rating = (lrand48 () >> 16) % n_labels;
    }
}

static void
test_to_from_identity_once (void)
{
  /* (alpha, logbeta) = from_x (alpha, logbeta, to_x (alpha, logbeta)) */

  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;

  double alpha[n_workers * n_labels * n_labels];
  double logbeta[n_images];
  double newalpha[n_workers * n_labels * n_labels];
  double newlogbeta[n_images];

  fill (alpha, drand48, n_workers * n_labels * n_labels);
  fill (newalpha, drand48, n_workers * n_labels * n_labels);
  fill (logbeta, drand48, n_images);
  fill (newlogbeta, drand48, n_images);

  size_t n_dims = n_workers * n_labels * n_labels + n_images;
  gsl_vector* x = gsl_vector_alloc (n_dims);

  to_x (alpha, n_workers, n_labels, logbeta, n_images, x);
  from_x (newalpha, n_workers, n_labels, newlogbeta, n_images, x);

  gsl_vector_free (x);

  assert (memcmp (alpha,
                  newalpha,
                  n_workers * n_labels * n_labels * sizeof (alpha[0])) == 0);

  assert (memcmp (logbeta, newlogbeta, n_images * sizeof (logbeta[0])) == 0);

}

static void
test_to_from_inverse (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    {
      test_to_from_identity_once ();
    }
}

static void
test_qfunc_f_once (void)
{
  /* qfunc (alpha, logbeta) = -qfunc_f (to_x (alpha, logbeta)) */

  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_ratings = ((lrand48 () >> 16) % 5) + 1;

  double alpha[n_workers * n_labels * n_labels];
  double logbeta[n_images];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  Rating ratings[n_ratings];
  NormalDistribution prioralpha = normal_distribution (drand48 (), 1.0 + drand48 ());
  NormalDistribution priorlogbeta = normal_distribution (drand48 (), 1.0 + drand48 ());

  fill (alpha, drand48, n_workers * n_labels * n_labels);
  fill (logbeta, drand48, n_images);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);
  fill_r (ratings, n_workers, n_labels, n_images, n_ratings);

  QfuncParams qfunc_params = { .n_workers = n_workers,
                               .n_labels = n_labels,
                               .n_images = n_images,
                               .pz = pz,
                               .priorz = priorz,
                               .ratings = ratings,
                               .n_ratings = n_ratings,
                               .prioralpha = &prioralpha.base,
                               .priorlogbeta = &priorlogbeta.base };

  size_t n_dims = n_workers * n_labels * n_labels + n_images;
  gsl_vector* x = gsl_vector_alloc (n_dims);

  to_x (alpha, n_workers, n_labels, logbeta, n_images, x);

  double q = qfunc (alpha, n_workers, n_labels, logbeta, n_images,
                    pz, priorz, ratings, n_ratings, 
                    &prioralpha.base, &priorlogbeta.base, 0);

  double qf = qfunc_f (x, &qfunc_params);

  assert (fabs (q + qf) <= 1e-8 * (1 + fabs (q) + fabs (qf)) ||
          (fprintf (stderr, "%g ?= %g\n", q, qf), 0));

  gsl_vector_free (x);
}

static void
test_qfunc_f (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    {
      test_qfunc_f_once ();
    }
}

static void
test_qfunc_df_once (void)
{
  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_ratings = ((lrand48 () >> 16) % 5) + 1;

  /* dqfunc (alpha, logbeta) = -qfunc_df (to_x (alpha, logbeta)) */

  double alpha[n_workers * n_labels * n_labels];
  double logbeta[n_images];
  double dalpha[n_workers * n_labels * n_labels];
  double dlogbeta[n_images];
  double newdalpha[n_workers * n_labels * n_labels];
  double newdlogbeta[n_images];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  Rating ratings[n_ratings];
  NormalDistribution prioralpha = normal_distribution (drand48 (), 1.0 + drand48 ());
  NormalDistribution priorlogbeta = normal_distribution (drand48 (), 1.0 + drand48 ());

  fill (alpha, drand48, n_workers * n_labels * n_labels);
  fill (logbeta, drand48, n_images);
  fill (dalpha, drand48, n_workers * n_labels * n_labels);
  fill (dlogbeta, drand48, n_images);
  fill (newdalpha, drand48, n_workers * n_labels * n_labels);
  fill (newdlogbeta, drand48, n_images);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);
  fill_r (ratings, n_workers, n_labels, n_images, n_ratings);

  QfuncParams qfunc_params = { .n_workers = n_workers,
                               .n_labels = n_labels,
                               .n_images = n_images,
                               .pz = pz,
                               .priorz = priorz,
                               .ratings = ratings,
                               .n_ratings = n_ratings,
                               .prioralpha = &prioralpha.base,
                               .priorlogbeta = &priorlogbeta.base };

  size_t n_dims = n_workers * n_labels * n_labels + n_images;
  gsl_vector* x = gsl_vector_alloc (n_dims);
  gsl_vector* g = gsl_vector_alloc (n_dims);

  to_x (alpha, n_workers, n_labels, logbeta, n_images, x);

  dqfunc (dalpha, dlogbeta, 
          alpha, n_workers, n_labels, logbeta, n_images,
          pz, ratings, n_ratings, 
          &prioralpha.base, &priorlogbeta.base, 0);

  qfunc_df (x, &qfunc_params, g);

  from_x (newdalpha, n_workers, n_labels, newdlogbeta, n_images, g);

  for (unsigned int i = 0; i < n_workers * n_labels * n_labels; ++i)
    {
      assert (   fabs (newdalpha[i] + dalpha[i]) 
              <= 1e-8 * (1 + fabs (newdalpha[i]) + fabs (dalpha[i])) ||
              (fprintf (stderr, "%u: %g ?= %g\n", i, newdalpha[i], dalpha[i]),
               0));
    }

  for (unsigned int i = 0; i < n_images; ++i)
    {
      assert (   fabs (newdlogbeta[i] + dlogbeta[i]) 
              <= 1e-8 * (1 + fabs (newdlogbeta[i]) + fabs (dlogbeta[i])) ||
              (fprintf (stderr, "%u: %g ?= %g\n", i, newdlogbeta[i], dlogbeta[i]),
               0));
    }

  gsl_vector_free (g);
  gsl_vector_free (x);
}

static void
test_qfunc_df (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    {
      test_qfunc_df_once ();
    }
}

static void
test_mstep_no_data_once (void)
{
  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;

  /* without data mstep should converge on prior maximum */

  double alpha[n_workers * n_labels * n_labels];
  double diagonal[n_workers * n_labels];
  double logbeta[n_images];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  NormalDistribution prioralpha = normal_distribution (drand48 (), 1.0 + drand48 ());
  NormalDistribution priorlogbeta = normal_distribution (drand48 (), 1.0 + drand48 ());

  fill (alpha, drand48, n_workers * n_labels * n_labels);
  fill (logbeta, drand48, n_images);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);

  for (unsigned int j = 0; j < n_workers; ++j)
    {
      for (unsigned int k = 0; k < n_labels; ++k)
        {
          diagonal[j * n_labels + k] = 
            alpha[j * n_labels * n_labels + k * n_labels + k];
        }
    }

  mstep (alpha, n_workers, n_labels, 
         logbeta, n_images, 
         pz, priorz, 
         NULL, 0, 
         &prioralpha.base, &priorlogbeta.base,
         0, 1e-8);

  for (unsigned int i = 0, x = 0; i < n_workers * n_labels * n_labels; ++i)
    {
      if (i % n_labels == (i / n_labels) % n_labels)
        {
          assert (   fabs (alpha[i] - diagonal[x]) <= 
                     1e-8 * (1 + fabs (alpha[i]) + fabs (diagonal[x])) ||
                  (fprintf (stderr, "%u: %g ?= %g\n", i, alpha[i], diagonal[x]),
                   0));

          ++x;
        }
      else
        {
          assert (   fabs (alpha[i] - prioralpha.mean) <= 
                     1e-8 * (1 + fabs (alpha[i]) + fabs (prioralpha.mean)) ||
                  (fprintf (stderr, "%u: %g ?= %g\n", i, alpha[i], prioralpha.mean),
                   0));
        }
    }

  for (unsigned int i = 0; i < n_images; ++i)
    {
      assert (   fabs (logbeta[i] - priorlogbeta.mean) <= 
                 1e-8 * (1 + fabs (logbeta[i]) + fabs (priorlogbeta.mean)) ||
              (fprintf (stderr, "%u: %g ?= %g\n", i, logbeta[i], priorlogbeta.mean),
               0));
    }
}

static void
test_mstep_no_data (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    {
      test_mstep_no_data_once ();
    }
}

static void
test_mstep_symmetric_no_data_once (void)
{
  unsigned int n_workers = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_labels = ((lrand48 () >> 16) % 5) + 1;
  unsigned int n_images = ((lrand48 () >> 16) % 5) + 1;

  /* without data mstep should converge on prior maximum */

  double alpha[n_workers * n_labels * n_labels];
  double alphacopy[n_workers * n_labels * n_labels];
  double logbeta[n_images];
  double pz[n_images * n_labels];
  double priorz[n_labels];
  NormalDistribution prioralpha = normal_distribution (drand48 (), 1.0 + drand48 ());
  NormalDistribution priorlogbeta = normal_distribution (drand48 (), 1.0 + drand48 ());

  fill (alpha, drand48, n_workers * n_labels * n_labels);
  memcpy (alphacopy, alpha, n_workers * n_labels * n_labels * sizeof (alpha[0]));
  fill (logbeta, drand48, n_images);
  fill (pz, drand48, n_images * n_labels);
  fill (priorz, drand48, n_labels);

  mstep (alpha, n_workers, n_labels, 
         logbeta, n_images, 
         pz, priorz, 
         NULL, 0, 
         &prioralpha.base, &priorlogbeta.base,
         1, 1e-8);

  for (unsigned int i = 0; i < n_workers * n_labels * n_labels; ++i)
    {
      if (i % n_labels <= (i / n_labels) % n_labels)
        {
          assert (   fabs (alpha[i] - alphacopy[i]) <= 
                     1e-8 * (1 + fabs (alpha[i]) + fabs (alphacopy[i])) ||
                  (fprintf (stderr, "%u: %g ?= %g\n", i, alpha[i], alphacopy[i]),
                   0));
        }
      else
        {
          assert (   fabs (alpha[i] - prioralpha.mean) <= 
                     1e-8 * (1 + fabs (alpha[i]) + fabs (prioralpha.mean)) ||
                  (fprintf (stderr, "%u: %g ?= %g\n", i, alpha[i], prioralpha.mean),
                   0));
        }
    }

  for (unsigned int i = 0; i < n_images; ++i)
    {
      assert (   fabs (logbeta[i] - priorlogbeta.mean) <= 
                 1e-8 * (1 + fabs (logbeta[i]) + fabs (priorlogbeta.mean)) ||
              (fprintf (stderr, "%u: %g ?= %g\n", i, logbeta[i], priorlogbeta.mean),
               0));
    }
}

static void
test_mstep_symmetric_no_data (void)
{
  for (unsigned int i = 0; i < 100; ++i)
    {
      test_mstep_symmetric_no_data_once ();
    }
}

typedef struct _MstepTestDatum MstepTestDatum ;
struct _MstepTestDatum 
{ 
  struct
    {
      double    alpha[4 * 5 * 5];
      double    logbeta[3];
      double    pz[3 * 5];
      double    priorz[5];
      Rating    ratings[8];
      double    alphamean;
      double    alphastddev;
      double    logbetamean;
      double    logbetastddev;
    }                   inputs;

  struct
    {
      double    newq;
      struct
        {
          double    alpha[4 * 5 * 5];
          double    logbeta[3];
        }       new;
    }                   desired_output;
};

static void
test_mstep (void)
{
  MstepTestDatum data[] = 
      {{{{0, 0.387338, 0.764869, 0.903349, 0.681943, 0.179307, 0, 0.885865,  0.559996, 0.716166, 0.844816, 0.421489, 0, 0.900841, 0.252985,  0.709759, 0.947329, 0.945259, 0, 0.664301, 0.848842, 0.342364,  0.289449, 0.970765, 0, 0, 0.317154, 0.308575, 0.430021, 0.437372,  0.929816, 0, 0.543706, 0.526673, 0.755429, 0.750509, 0.657841, 0,  0.966677, 0.0392629, 0.905693, 0.236352, 0.0658357, 0, 0.786278,  0.195934, 0.289023, 0.120577, 0.121976, 0, 0, 0.347092, 0.946659,  0.831128, 0.151212, 0.0299387, 0, 0.638084, 0.401107, 0.713839,  0.100122, 0.0943779, 0, 0.874434, 0.95841, 0.349613, 0.436537,  0.907757, 0, 0.919147, 0.44392, 0.200185, 0.841921, 0.13287, 0, 0,  0.247985, 0.911162, 0.721344, 0.0108934, 0.900893, 0, 0.964504,  0.890217, 0.859682, 0.870954, 0.32642, 0, 0.48911, 0.145842,  0.770832, 0.232042, 0.614676, 0, 0.187432, 0.421219, 0.795505,  0.706918, 0.268285, 0}, {0.977299, 0.59532, 0.864997}, {0.135415,  0.729313, 0.684157, 0.143652, 0.124521, 0.828421, 0.719654,  0.253436, 0.26484, 0.957466, 0.393234, 0.764326, 0.118998,  0.186635, 0.161193}, {0.14965, 0.931566, 0.765416, 0.365688,  0.442732}, {{3, 2, 4}, {2, 0, 4}, {0, 1, 4}, {3, 2, 0}, {1, 0,  2}, {0, 1, 2}, {1, 0, 1}, {0, 1, 2}}, 0.663281, 0.788117,  0.770368,  0.577735}, {-104.951, {{0., 0.752403, -0.209376, 0.752403,  0.0261426, 0.749057, 0., -0.17497, 0.749057, 0.0637509, 0.743776,  0.743776, 0., 0.743776, 0.387341, 0.723005, 0.723005, 0.0809844,  0., 0.323288, 0.802148, 0.802148, -0.096904, 0.802148, 0., 0.,  0.499083, 0.499083, 0.700804, 0.700804, 0.805716, 0., 0.190471,  0.805716, 0.805716, 0.799959, 0.202167, 0., 0.799959, 0.799959,  0.702649, 0.490335, 0.490335, 0., 0.702649, 0.698296, 0.51086,  0.51086, 0.698296, 0., 0., 0.683689, 0.683689, 0.683689,  0.470544, 0.730657, 0., 0.730657, 0.730657, -0.0323198, 0.728589,  0.728589, 0., 0.728589, -0.00828677, 0.684781, 0.684781,  0.684781, 0., 0.459812, 0.68347, 0.68347, 0.68347, 0.68347, 0.,  0., 0.763247, 0.763247, 0.763247, 0.366946, 0.156852, 0.,  0.795002, 0.795002, 0.156852, 0.549387, 0.697795, 0., 0.697795,  0.549387, 0.492021, 0.713913, 0.713913, 0., 0.492021, 0.516611,  0.710579, 0.710579, 0.710579, 0.}, {1.00254, 1.27925,  0.782497}}}}, {{{0, 0.527866, 0.0588038, 0.0862103, 0.434083,  0.403345, 0, 0.230383, 0.366556, 0.180647, 0.138505, 0.272917, 0,  0.973322, 0.416321, 0.0195074, 0.0862823, 0.812129, 0, 0.266671,  0.0879418, 0.320866, 0.446442, 0.823938, 0, 0, 0.424661, 0.532749,  0.676074, 0.246203, 0.896795, 0, 0.473945, 0.589864, 0.81212,  0.49345, 0.243562, 0, 0.223307, 0.631473, 0.354945, 0.970645,  0.249985, 0, 0.215152, 0.335438, 0.884363, 0.437856, 0.948482, 0,  0, 0.247496, 0.563496, 0.991414, 0.124543, 0.822835, 0, 0.0307473,  0.31534, 0.87834, 0.92604, 0.556802, 0, 0.725476, 0.0662198,  0.432591, 0.31324, 0.502169, 0, 0.434746, 0.0776457, 0.342595,  0.252184, 0.219594, 0, 0, 0.742208, 0.458232, 0.814328, 0.271112,  0.494712, 0, 0.894735, 0.822913, 0.146569, 0.671877, 0.863988, 0,  0.507573, 0.268229, 0.745837, 0.307186, 0.782097, 0, 0.202009,  0.313246, 0.993946, 0.279928, 0.767262, 0}, {0.235601, 0.651351,  0.027744}, {0.547668, 0.493392, 0.19312, 0.213417, 0.276556,  0.99868, 0.298384, 0.390503, 0.129987, 0.326802, 0.434396,  0.88293, 0.861758, 0.580966, 0.12721}, {0.100834, 0.65975,  0.267719, 0.133264,  0.820906}, {{0, 2, 1}, {1, 2, 1}, {2, 1, 2}, {1, 1, 0}, {1, 0,  0}, {0, 1, 4}, {1, 1, 1}, {2, 2, 0}}, 0.892487, 0.0321188,  0.481913,  0.793162}, {-7511.98, {{0., 0.892279, 0.892706, 0.892706, 0.891757,  0.892669, 0., 0.892669, 0.892669, 0.892386, 0.89268, 0.891833,  0., 0.89268, 0.892309, 0.892597, 0.892026, 0.892597, 0.,  0.892473, 0.892556, 0.892431, 0.892556, 0.892556, 0., 0.,  0.891586, 0.892961, 0.892961, 0.892961, 0.891836, 0., 0.892808,  0.892808, 0.892808, 0.892135, 0.891558, 0., 0.892775, 0.892775,  0.892236, 0.891963, 0.892657, 0., 0.892657, 0.891962, 0.892223,  0.892658, 0.892658, 0., 0., 0.892706, 0.891757, 0.892706,  0.892706, 0.891802, 0., 0.892386, 0.892669, 0.892669, 0.891833,  0.89268, 0., 0.89268, 0.89268, 0.892026, 0.892597, 0.892473, 0.,  0.892597, 0.892431, 0.892556, 0.892246, 0.892556, 0., 0.,  0.892487, 0.892487, 0.892487, 0.892487, 0.892487, 0., 0.892487,  0.892487, 0.892487, 0.892487, 0.892487, 0., 0.892487, 0.892487,  0.892487, 0.892487, 0.892487, 0., 0.892487, 0.892487, 0.892487,  0.892487, 0.892487,  0.}, {0.302086, -0.0823632, -0.0487221}}}}, {{{0, 0.344819,  0.538726, 0.288793, 0.579746, 0.0682628, 0, 0.540047, 0.990409,  0.189242, 0.938276, 0.213244, 0, 0.556012, 0.306312, 0.0765173,  0.632278, 0.428802, 0, 0.205478, 0.416768, 0.364559, 0.295538,  0.384572, 0, 0, 0.524281, 0.33244, 0.813625, 0.59141, 0.179462, 0,  0.793714, 0.524832, 0.0116646, 0.111199, 0.253667, 0, 0.534423,  0.822422, 0.172923, 0.0404232, 0.97841, 0, 0.51611, 0.0964061,  0.408145, 0.549608, 0.310632, 0, 0, 0.679638, 0.0435855, 0.254071,  0.926059, 0.155358, 0, 0.711145, 0.440446, 0.334649, 0.975896,  0.917431, 0, 0.915614, 0.322985, 0.864697, 0.663764, 0.381191, 0,  0.500562, 0.691773, 0.623341, 0.402781, 0.984452, 0, 0, 0.595367,  0.215196, 0.853172, 0.67382, 0.915729, 0, 0.171611, 0.599102,  0.747761, 0.760371, 0.460465, 0, 0.158656, 0.413112, 0.784476,  0.543034, 0.243042, 0, 0.0901273, 0.919779, 0.87927, 0.86185,  0.589565, 0}, {0.228005, 0.255929, 0.45907}, {0.605113, 0.632638,  0.0407332, 0.605897, 0.931292, 0.716909, 0.869123, 0.00679579,  0.183531, 0.956538, 0.408657, 0.84814, 0.770419, 0.172062,  0.865623}, {0.605098, 0.680292, 0.252284, 0.986353,  0.743248}, {{2, 1, 1}, {1, 1, 3}, {3, 1, 1}, {0, 2, 4}, {2, 0,  3}, {1, 2, 4}, {2, 2, 3}, {3, 1, 4}}, 0.0907271, 0.0242782,  0.730424,  0.284178}, {100.733, {{0., 0.0908229, 0.0908229, 0.0908229,  0.0903242, 0.0909258, 0., 0.0909258, 0.0909258, 0.089891,  0.0909076, 0.0909076, 0., 0.0909076, 0.0899676, 0.0907674,  0.0907674, 0.0907674, 0., 0.0905574, 0.09093, 0.09093, 0.09093,  0.09093, 0., 0., 0.0909922, 0.0909922, 0.0901101, 0.0904936,  0.0911309, 0., 0.0911309, 0.0900619, 0.0900965, 0.0909092,  0.0909092, 0., 0.0909009, 0.0899692, 0.0908108, 0.0908108,  0.0908108, 0., 0.0906008, 0.0911558, 0.0911558, 0.0911558,  0.0899793, 0., 0., 0.0902523, 0.0911341, 0.0898961, 0.0911341,  0.0912794, 0., 0.0912794, 0.0894723, 0.0912794, 0.0909188,  0.0909104, 0., 0.089929, 0.0909188, 0.0909531, 0.0907273,  0.0909531, 0., 0.0909531, 0.0913738, 0.0901979, 0.0913738,  0.0891809, 0., 0., 0.0901838, 0.0910657, 0.0910657, 0.0901838,  0.0911377, 0., 0.0911377, 0.0911377, 0.0900687, 0.0907303,  0.090722, 0., 0.0907303, 0.090722, 0.0908138, 0.090588,  0.0908138, 0., 0.090588, 0.0911789, 0.0900025, 0.0911789,  0.0911789, 0.}, {0.730163, 0.736466, 0.727969}}}}}; 

  for (unsigned int i = 0; i < sizeof (data) / sizeof (data[0]); ++i)
    {
      NormalDistribution prior_alpha = 
        normal_distribution (data[i].inputs.alphamean, 
                             data[i].inputs.alphastddev);
      NormalDistribution prior_logbeta = 
        normal_distribution (data[i].inputs.logbetamean, 
                             data[i].inputs.logbetastddev);
      double diagonal[4 * 5];

      for (unsigned int j = 0; j < 4; ++j)
        {
          for (unsigned int k = 0; k < 5; ++k)
            {
              diagonal[j * 5 + k] = drand48 ();
              data[i].inputs.alpha[j * 5 * 5 + k * 5 + k] = diagonal[j * 5 + k];
            }
        }

      double q = mstep (data[i].inputs.alpha,
                        4,
                        5,
                        data[i].inputs.logbeta,
                        3,
                        data[i].inputs.pz,
                        data[i].inputs.priorz,
                        data[i].inputs.ratings,
                        8,
                        &prior_alpha.base,
                        &prior_logbeta.base,
                        0,
                        1e-8);

      assert (fabs (q - data[i].desired_output.newq) <=
              1e-4 * (1 + fabs (q) + fabs (data[i].desired_output.newq)) ||
              ( fprintf (stderr, "%g ?= %g\n", q, data[i].desired_output.newq),
                0 ));
      
      for (unsigned int j = 0, y = 0; j < 4*5*5; ++j)
        {
          if (j % 5 == (j / 5) % 5)
            {
              assert (fabs (data[i].inputs.alpha[j] - diagonal[y]) <=
                      1e-8 * (1 + fabs (data[i].inputs.alpha[j]) + fabs (diagonal[y])) ||
                      ( fprintf (stderr, "%u %u: %g ?= %g\n", i, j, data[i].inputs.alpha[j], diagonal[y]),
                        0 ));

              ++y;
            }
          else
            {
              assert (fabs (data[i].inputs.alpha[j] - data[i].desired_output.new.alpha[j]) <=
                      5e-3 * (1 + fabs (data[i].inputs.alpha[j]) + fabs (data[i].desired_output.new.alpha[j])) ||
                      ( fprintf (stderr, "%u %u: %g ?= %g\n", i, j, data[i].inputs.alpha[j], data[i].desired_output.new.alpha[j]),
                        0 ));
            }
        }
    }
}

int 
main (void)
{
  srand (69);

  test_to_from_inverse ();
  test_qfunc_f ();
  test_qfunc_df ();
  test_mstep_no_data ();
  test_mstep_symmetric_no_data ();
  test_mstep ();

  /* TODO: copy longer test vectors in */

  return 0;
}
